import os
import requests
from dotenv import load_dotenv
from requests.auth import HTTPBasicAuth
from urllib3.exceptions import InsecureRequestWarning

# 加载环境变量
load_dotenv()

# 禁用 SSL 验证警告
import urllib3
urllib3.disable_warnings(category=InsecureRequestWarning)

class Ollama:

    def __init__(self):
        self.__username = os.getenv("CADDY_USERNAME")
        self.__password = os.getenv("CADDY_PASSWORD")
        self.__endpoint = os.getenv("OLLAMA_ENDPOINT")
        if not self.__endpoint:
            raise ValueError("OLLAMA_ENDPOINT is not set in the environment variables.")
        if not self.__username or not self.__password:
            raise ValueError("CADDY_USERNAME or CADDY_PASSWORD is not set in the environment variables.")

    def tags(self):
        url = f"{self.__endpoint}/api/tags"
        return requests.get(url, auth=HTTPBasicAuth(self.__username, self.__password), verify=False)

    def ps(self):
        url = f"{self.__endpoint}/api/ps"
        return requests.get(url, auth=HTTPBasicAuth(self.__username, self.__password), verify=False)

    def generate(self, model, prompt, **kwargs):
        url = f"{self.__endpoint}/api/generate"
        payload = {
            "model": model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": kwargs.get("temperature", 0.7),
                "top_p": kwargs.get("top_p", 0.9),
                "max_tokens": kwargs.get("max_tokens", 512),
                "num_gpu_layers": 40  # GPU加速层数
            }
        }
        return requests.post(url, json=payload, auth=HTTPBasicAuth(self.__username, self.__password), verify=False)

class Deepseek:

    def __init__(self):
        self.__base_url = "https://api.deepseek.com/v1"
        self.__api_key = os.getenv("DEEPSEEK_OFFICIAL_KEY")
        if not self.__api_key:
            raise ValueError("DEEPSEEK_OFFICIAL_KEY is not set in the environment variables.")
        self.__headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.__api_key}"
        }

    def chat_completions(self, prompt: str, model: str = "deepseek-chat", **kwargs):
        url = f"{self.__base_url}/chat/completions"
        payload = {
            "model": model,
            "messages": [
                {"role": "system", "content": prompt}
            ],
            "temperature": kwargs.get("temperature", 0.7),
            "top_p": kwargs.get("top_p", 1.0),
            "max_tokens": kwargs.get("max_tokens", 1024),
            "stream": kwargs.get("stream", False)
        }
        return requests.post(url, json=payload, headers=self.__headers, timeout=15)

    def models(self):
        url = f"{self.__base_url}/models"
        return requests.get(url, headers=self.__headers, timeout=10)

    def usage(self):
        url = f"{self.__base_url}/usage"
        return requests.get(url, headers=self.__headers, timeout=10)